package com.rrd.hive.udf;

import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDF;
import org.apache.hadoop.hive.ql.udf.xml.UDFXPathUtil;
import org.apache.hadoop.hive.ql.exec.Description;

import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector.Category;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorConverters;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorConverters.Converter;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorUtils;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorUtils.PrimitiveGrouping;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;

import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerException;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;
import java.io.StringWriter;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.LinkedHashMap;
import java.util.List;

@Description(
        name = "xml_to_map",
        value = "_FUNC_(text) - Creates a map by xml string ",
        extended = "parse xml string to map"
)
public class GenericUDFXmlToMap extends GenericUDF {
    private final LinkedHashMap<Object, Object> ret = new LinkedHashMap();
    private transient Converter soi_text, soi_text1 = null, soi_text2 = null;
    private UDFXPathUtil xpath;

    public GenericUDFXmlToMap() {
    }

    public ObjectInspector initialize(ObjectInspector[] arguments) throws UDFArgumentException {
        for(int idx = 0; idx < Math.min(arguments.length, 3); ++idx) {
            if (arguments[idx].getCategory() != Category.PRIMITIVE || PrimitiveObjectInspectorUtils.getPrimitiveGrouping(((PrimitiveObjectInspector)arguments[idx]).getPrimitiveCategory()) != PrimitiveGrouping.STRING_GROUP) {
                throw new UDFArgumentException("All argument should be string/character type");
            }
        }

        this.soi_text = ObjectInspectorConverters.getConverter(arguments[0], PrimitiveObjectInspectorFactory.javaStringObjectInspector);
        if (arguments.length >= 2) {
            this.soi_text1 = ObjectInspectorConverters.getConverter(arguments[1], PrimitiveObjectInspectorFactory.javaStringObjectInspector);
        }
        if (arguments.length >= 3) {
            this.soi_text2 = ObjectInspectorConverters.getConverter(arguments[2], PrimitiveObjectInspectorFactory.javaStringObjectInspector);
        }
        this.xpath = new UDFXPathUtil();
        return ObjectInspectorFactory.getStandardMapObjectInspector(PrimitiveObjectInspectorFactory.javaStringObjectInspector, PrimitiveObjectInspectorFactory.javaStringObjectInspector);
    }

    private static String toString(Node n) {
        String value = null;
        if (n.getChildNodes().getLength() == 1 && n.getFirstChild().getNodeType() == 3) {
            value = n.getFirstChild().getNodeValue();
        }
        if (value == null) {
            try{
                Transformer transformer = TransformerFactory.newInstance().newTransformer();
                StreamResult result = new StreamResult(new StringWriter());
                DOMSource source = new DOMSource(n);
                transformer.transform(source, result);
                value = result.getWriter().toString().replace("<?xml version=\"1.0\" encoding=\"UTF-8\"?>", "");
            }catch (TransformerException e) {
                e.printStackTrace();
            }
        }
        return value;
    }

    public Object evaluate(DeferredObject[] arguments) throws HiveException {
        this.ret.clear();
        String xml = (String)this.soi_text.convert(arguments[0].get());
        String path = this.soi_text1 == null ? "" : (String)this.soi_text1.convert(arguments[1].get());
        String exclude = this.soi_text2 == null ? "" : (String)this.soi_text2.convert(arguments[2].get());
        if (xml == null) {
            return null;
        }
        xml = xml.trim();
        if (path.equals("")) {
            path = "xml";
            xml = "<xml>" + xml + "</xml>";
        }
        List<String> excludes = new ArrayList<>();
        if (!exclude.equals("")){
            excludes = Arrays.asList(exclude.split(","));
        }
        Node node = xpath.evalNode(xml, path);
        if (node != null){
            NodeList childNodes = node.getChildNodes();
            for (int i = 0; i < childNodes.getLength(); i++) {
                Node childNode = childNodes.item(i);
                String name = childNode.getNodeName();

                if (excludes.size() == 0 || !excludes.contains(name)){
                    this.ret.put(name, toString(childNode));
                }
            }
        }
        return this.ret;

    }

    public String getDisplayString(String[] children) {
        assert children.length <= 3;
        return this.getStandardDisplayString("xml_to_map", children, ",");
    }

}
